import triton
import triton.language as tl

from torch._inductor.runtime import triton_helpers
from torch._inductor.runtime.triton_helpers import math as tl_math

import torch
import torch_npu
from .util import triton_helper_compile, compiled_kernel_list

# triton_per_fused__safe_softmax_add_eq_masked_fill_4
@triton.jit
def fused_kernel10(in_ptr0, in_ptr1, in_ptr2, out_ptr3, xnumel, rnumel):
    xnumel = 32
    XBLOCK: tl.constexpr = 1
    rnumel = 10
    RBLOCK: tl.constexpr = 16
    xoffset = tl.program_id(0) * XBLOCK
    xindex = tl.full([1], xoffset, tl.int32)
    xmask = tl.full([RBLOCK], True, tl.int1)
    rindex = tl.arange(0, RBLOCK)[:]
    roffset = 0
    rmask = rindex < rnumel
    r1 = rindex
    x0 = xindex
    tmp0 = tl.load(in_ptr0 + (r1 + (10*x0)), rmask, other=0.0)
    tmp1 = tl.load(in_ptr1 + (0))
    tmp2 = tl.broadcast_to(tmp1, [RBLOCK])
    tmp10 = tl.load(in_ptr2 + (r1), rmask, eviction_policy='evict_last', other=0.0)
    tmp3 = tmp2.to(tl.float32)
    tmp4 = r1
    tmp5 = tmp4.to(tl.float32)
    tmp6 = tmp5 > tmp3
    tmp7 = tmp6.to(tl.float32)
    tmp8 = -3.4028234663852886e+38
    tmp9 = tmp8 * tmp7
    tmp11 = tmp10.to(tl.float32)
    tmp12 = tmp9 + tmp11
    tmp13 = 0.0
    tmp14 = tmp12 == tmp13
    tmp15 = tl.where(tmp14, tmp8, tmp9)
    tmp16 = tmp0 + tmp15
    tmp17 = float("-inf")
    tmp18 = tmp16 == tmp17
    tmp19 = tmp18 == 0
    tmp20 = tmp19.to(tl.int32)
    tmp21 = (tmp20 != 0)
    tmp22 = tl.broadcast_to(tmp21, [RBLOCK])
    tmp24 = tl.where(rmask, tmp22, 0)
    tmp25 = triton_helpers.promote_to_tensor(triton_helpers.any(tmp24, 0))
    tmp26 = tl.broadcast_to(tmp16, [RBLOCK])
    tmp28 = tl.where(rmask, tmp26, float("-inf"))
    tmp29 = triton_helpers.promote_to_tensor(triton_helpers.max2(tmp28, 0))
    tmp30 = tmp16 - tmp29
    tmp31 = tl_math.exp(tmp30)
    tmp32 = tl.broadcast_to(tmp31, [RBLOCK])
    tmp34 = tl.where(rmask, tmp32, 0)
    tmp35 = triton_helpers.promote_to_tensor(tl.sum(tmp34, 0))
    tmp36 = tmp25 == 0
    tmp37 = tmp31 / tmp35
    tmp38 = tl.where(tmp36, tmp13, tmp37)
    tl.store(out_ptr3 + (r1 + (10*x0)), tmp38, rmask)

def test_compile_for_fused_kernel10():
    signature = {0: '*fp32', 1: '*i32', 2: '*i32', 3: '*fp32', 4: 'i32', 5: 'i32'}
    triton_helper_compile(func=fused_kernel10, signature=signature)

compiled_kernel_list["fused_kernel10"] = test_compile_for_fused_kernel10